import chromadb
from config import CHROMA_PERSIST_DIR, COLLECTION_NAME


import chromadb
import logging

from config import CHROMA_PERSISTENT_DATA_PATH, CHROMA_SERVER_HOST, CHROMA_SERVER_POST, \
    CHROMA_CONTRACT_NAME


class Chroma(object):
    def __init__(self):
        self.client = self.get_httpclient()
        self.collection = self.client.get_or_create_collection(name=CHROMA_CONTRACT_NAME,
                                                               metadata={"hnsw:space": "cosine"}) \
            if self.client else None

    @classmethod
    def get_client(cls):
        """
        Client() 基于内存存储，数据不会持久化存储，不建议用于生产环境
        """
        client = chromadb.Client()
        return client

    @classmethod
    def get_persistent_client(cls):
        """
        PersistentClient() 基于硬盘存储，可以持久化存储，会把数据存到指定路径下，但不建议用于生产环境
        来保存和加载本地机器上的数据，启动时加载(如果存在)
        """
        try:
            client = chromadb.PersistentClient(path=CHROMA_PERSISTENT_DATA_PATH)
            return True
        except Exception as ex:
            logging.info("保存加载本地机器上的数据失败，原因是：%s" % str(ex))
            return False

    @classmethod
    def get_httpclient(cls):
        """
        HttpClient() 连接到Chroma服务器进行存储，可多台客户端连接同一个服务器，可用于生产环境
        """
        try:
            client = chromadb.HttpClient(
                host=CHROMA_SERVER_HOST,
                port=CHROMA_SERVER_POST
            )
            Chroma.get_persistent_client()
            return client
        except Exception as ex:
            print(f"连接chroma服务器失败，原因是：{str(ex)}")
            logging.info(f"连接chroma服务器失败，原因是：{str(ex)}")
            return None

    def chromadb(self):
        return self.collection
